{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0be845da",
   "metadata": {},
   "source": [
    "## SAM 能力验证"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33681dd1",
   "metadata": {},
   "source": [
    "导入相关的包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b28288",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29bc90d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_mask(mask, ax, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        color = np.array([30/255, 144/255, 255/255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "    \n",
    "def show_points(coords, labels, ax, marker_size=375):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n",
    "    \n",
    "def show_box(box, ax):\n",
    "    x0, y0 = box[0], box[1]\n",
    "    w, h = box[2] - box[0], box[3] - box[1]\n",
    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23842fb2",
   "metadata": {},
   "source": [
    "## 测试图片"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c2e4f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = cv2.imread('images/truck.jpg')\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ba2c244",
   "metadata": {},
   "outputs": [],
   "source": [
    "image.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e30125fd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "plt.imshow(image)\n",
    "plt.axis('on')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b228b8",
   "metadata": {},
   "source": [
    "## 使用SAM分割"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e28150b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from segment_anything import sam_model_registry, SamPredictor\n",
    "# 模型权重文件\n",
    "sam_checkpoint = \"sam_vit_h_4b8939.pth\"\n",
    "# 模型类型\n",
    "model_type = \"vit_h\"\n",
    "# 使用GPU\n",
    "device = \"cuda\"\n",
    "# 加载模型\n",
    "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "# 将模型放到GPU上\n",
    "sam.to(device=device)\n",
    "# 创建预测器\n",
    "predictor = SamPredictor(sam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d95d48dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用`SamPredictor.set_image`生成embedding。 `SamPredictor`将记住这个值，并将其用于后续的掩膜预测。\n",
    "predictor.set_image(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c69570c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 选择一个点 (x, y) 作为输入\n",
    "input_point = np.array([[500, 375]])\n",
    "# 标注这个点的类别，1 表示foreground，即前景点，0 表示background，即背景点\n",
    "input_label = np.array([1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a91ba973",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 绘制检查一下\n",
    "plt.figure(figsize=(10,10))\n",
    "plt.imshow(image)\n",
    "show_points(input_point, input_label, plt.gca())\n",
    "plt.axis('on')\n",
    "plt.show()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5373fd68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用预训练模型进行预测，参数：点坐标，点标签，是否输出多分支mask，输出：多分支mask，分数，logits\n",
    "# 当multimask_output=True时，输出3个分支的mask，为False时，输出1个分支的mask。未解决prompt提示模糊问题，建议使用True，然后通过scores进行筛选\n",
    "masks, scores, logits = predictor.predict(\n",
    "    point_coords=input_point,\n",
    "    point_labels=input_label,\n",
    "    multimask_output=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe5c85f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#  (number_of_masks) x H x W\n",
    "masks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8062ea1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# masks[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4d4833e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 分别显示三个mask\n",
    "plt.figure(figsize=(15,5))\n",
    "for i in range(3):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    plt.imshow(masks[i], cmap='gray')\n",
    "    plt.axis('off')\n",
    "    plt.title('mask'+str(i+1))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e812a7aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 三个mask的预测分数\n",
    "scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcf09e51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# logit是没有经过softmax之前的值\n",
    "logits.shape "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8758da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb686d1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 分别显示三个logits，logits表示每个像素点属于哪个类别的概率，logits越大，表示该像素点属于该类别的概率越大\n",
    "plt.figure(figsize=(15,5))\n",
    "for i in range(3):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    plt.imshow(logits[i], cmap='plasma') # plasma表示热力图，越小越紫，越大越黄\n",
    "    plt.axis('off')\n",
    "    plt.title('logit'+str(i+1))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c227a6",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# 将结果绘制在图像上\n",
    "plt.figure(figsize=(20,10))\n",
    "for i, (mask, score) in enumerate(zip(masks, scores)):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    plt.imshow(image)\n",
    "    show_mask(mask, plt.gca())\n",
    "    show_points(input_point, input_label, plt.gca())\n",
    "    plt.title(f\"Mask {i+1}, Score: {score:.3f}\", fontsize=18)\n",
    "    plt.axis('off')\n",
    "plt.show()\n",
    "\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e39b67e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_copy = image.copy()\n",
    "# 使用图片数组的方式贴图\n",
    "seg = masks[0]\n",
    "# 若mask为true，则将img_copy中的像素置为半透明\n",
    "alpha = 0.7\n",
    "color = (255,0,255)\n",
    "# 意思是所有mask为true的地方 一半像素来自于原始像素，另一半来自于 color\n",
    "img_copy[seg > 0,0] = img_copy[seg > 0,0] * alpha + color[0] * (1-alpha) #b\n",
    "img_copy[seg > 0,1] = img_copy[seg > 0,1] * alpha + color[1] * (1-alpha) #g\n",
    "img_copy[seg > 0,2] = img_copy[seg > 0,2] * alpha + color[2] * (1-alpha) #r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f35b7eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 显示\n",
    "plt.imshow(img_copy)\n",
    "plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e33d94ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20,10))\n",
    "alpha = 0.5\n",
    "colors  = [(255,0,255), (0,255,255), (0,255,0)]\n",
    "for i, (mask, score) in enumerate(zip(masks, scores)):\n",
    "    plt.subplot(1,3,i+1)\n",
    "    img_copy = image.copy()\n",
    "    # 使用图片数据的方式贴图\n",
    "    seg = masks[i]\n",
    "    # 若mask为true，则将img_copy中的像素置为半透明\n",
    "    color =colors[i]\n",
    "    img_copy[seg > 0,0] = img_copy[seg > 0,0] * alpha + color[0] * (1-alpha)\n",
    "    img_copy[seg > 0,1] = img_copy[seg > 0,1] * alpha + color[1] * (1-alpha)\n",
    "    img_copy[seg > 0,2] = img_copy[seg > 0,2] * alpha + color[2] * (1-alpha)\n",
    "\n",
    "    plt.imshow(img_copy)\n",
    "    \n",
    "    plt.title(f\"Mask {i+1}, Score: {score:.3f}\", fontsize=18)\n",
    "    plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fa31f7c",
   "metadata": {},
   "source": [
    "## 使用更多点确定目标"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6923b94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输入的2个点坐标\n",
    "input_point = np.array([[500, 375], [1125, 625]])\n",
    "# 都是1，表示都是正样本\n",
    "input_label = np.array([1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62d1711d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If available, a mask from a previous iteration can also be supplied to the model to aid in prediction.\n",
    "# 如果可用，还可以向模型提供来自上一次迭代的掩码以帮助预测。\n",
    "\n",
    "mask_input = logits[np.argmax(scores), :, :]  # 选择最大分值的那个类别的logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d98f96a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用预训练模型进行预测，参数：点坐标，点标签，输入mask，是否输出多个mask（在多点确定一个物体时，可以将此值设为False）\n",
    "\n",
    "masks, _, _ = predictor.predict(\n",
    "    point_coords=input_point,\n",
    "    point_labels=input_label,\n",
    "    mask_input=mask_input[None, :, :],\n",
    "    multimask_output=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce8b82f",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e06d5c8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "plt.imshow(image)\n",
    "show_mask(masks, plt.gca())\n",
    "show_points(input_point, input_label, plt.gca())\n",
    "plt.axis('off')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a196f68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如果只想要玻璃而不要车身，可以将第二个点的label改为0\n",
    "input_point = np.array([[500, 375], [1125, 625]])\n",
    "input_label = np.array([1, 0])\n",
    "\n",
    "mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81a52282",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks, _, _ = predictor.predict(\n",
    "    point_coords=input_point,\n",
    "    point_labels=input_label,\n",
    "    mask_input=mask_input[None, :, :],\n",
    "    multimask_output=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfca709f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(image)\n",
    "show_mask(masks, plt.gca())\n",
    "show_points(input_point, input_label, plt.gca())\n",
    "plt.axis('off')\n",
    "plt.show() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41e2d5a9",
   "metadata": {},
   "source": [
    "## 使用矩形框作为提示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ea92a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 矩形框坐标：l, t, r, b\n",
    "input_box = np.array([425, 600, 700, 875])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b35a8814",
   "metadata": {},
   "outputs": [],
   "source": [
    "# box：shape (n, 4) \n",
    "masks, _, _ = predictor.predict(\n",
    "    point_coords=None,\n",
    "    point_labels=None,\n",
    "    box=input_box[None, :], \n",
    "    multimask_output=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "984b79c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(image)\n",
    "show_mask(masks[0], plt.gca())\n",
    "show_box(input_box, plt.gca())\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1ed9f0a",
   "metadata": {},
   "source": [
    "## 点与矩形框结合提示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90e2e547",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 这里使用电鱼矩形框结合，只选择车轮胎（将车轮胎中心点的label设为0，移除选区）\n",
    "input_box = np.array([425, 600, 700, 875])\n",
    "input_point = np.array([[575, 750]])\n",
    "input_label = np.array([0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6956d8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "masks, _, _ = predictor.predict(\n",
    "    point_coords=input_point,\n",
    "    point_labels=input_label,\n",
    "    box=input_box,\n",
    "    multimask_output=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e13088a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(image)\n",
    "show_mask(masks[0], plt.gca())\n",
    "show_box(input_box, plt.gca())\n",
    "show_points(input_point, input_label, plt.gca())\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f27107",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c97473b1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
